
const fs = require('fs');
const tf = require('@tensorflow/tfjs-node');

const img2x = (buffer) => {
    return tf.tidy(() => { // 最佳实践，使用该方法会去除掉一些中间变量
        // ts 需要将所有的数据转换为 ts 格式
        const imgTs = tf.node.decodeImage(new Uint8Array(buffer));
        // 为预训练模型准备数据，该模型只接收尺寸为 224*224 大小的图片
        const imgTsResized = tf.image.resizeBilinear(imgTs, [224, 224]);
        // toFloat().sub(255 / 2).div(255 / 2) 为归一化操作，也是预训练模型要求的数据格式
        // 224, 224 为图片的尺寸，3 表示 RGB 彩色图片
        return imgTsResized.toFloat().sub(255 / 2).div(255 / 2).reshape([1, 224, 224, 3]);
    })
}

const getData = async (trainDir, outputDir) => {
    console.log('It is getData');

    const classes = fs.readdirSync(trainDir).filter(item => !item.includes('.'));
    fs.writeFileSync(`${outputDir}/classes.json`, JSON.stringify(classes));

    const inputs = [];
    const labels = [];
    classes.forEach((dir, dirIndex) => {
        fs.readdirSync(`${trainDir}/${dir}`).
        filter(item => item.match(/jpeg$/)).forEach(filename => {
            const imgPath = `${trainDir}/${dir}/${filename}`
            console.log(imgPath);
            const buffer = fs.readFileSync(imgPath);
            const x = img2x(buffer);

            inputs.push(x);
            labels.push(dirIndex);
        });
    });

    const xs = tf.concat(inputs);
    const ys = tf.tensor(labels);

    // 返回数据
    return {
        xs,
        ys,
        classes,
    }
}

module.exports = getData;